package com.tianji.aigc.service.impl;

import com.tianji.aigc.dto.ChatDTO;
import com.tianji.aigc.entity.ChatSession;
import com.tianji.aigc.enums.ChatEventTypeEnum;
import com.tianji.aigc.service.IChatService;
import com.tianji.aigc.service.IChatSessionService;
import com.tianji.aigc.vo.ChatEventVO;
import com.tianji.common.exceptions.BizIllegalException;
import org.springframework.ai.chat.client.ChatClient;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.stereotype.Service;
import reactor.core.publisher.Flux;

import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;

@Service
public class ChatServiceImpl implements IChatService {
    @Autowired
    private ChatClient chatClient;
    @Autowired
    private IChatSessionService chatSessionService;
    //存储大模型的生成状态，这里采用ConcurrentHashMap是确保线程安全，如果分布式环境可考虑用redis
    private static Map<String,Boolean> SESSION_MAP=new ConcurrentHashMap<String,Boolean>();
    @Override
    public Flux<ChatEventVO> chat(ChatDTO dto) {
        //判断会话是否还存在
        ChatSession chatSession = chatSessionService.lambdaQuery().eq(ChatSession::getSessionId, dto.getSessionId()).one();
        if (chatSession == null) {
            throw new BizIllegalException("会话不存在");
        }

        return chatClient.prompt()
                .user(dto.getQuestion())
                .stream()
                .content()
                .doFirst(()->SESSION_MAP.put(dto.getSessionId(), true))
                .doOnComplete(()->SESSION_MAP.remove(dto.getSessionId()))
                .doOnError(x->SESSION_MAP.remove(dto.getSessionId()))
                .takeWhile(takewhile->SESSION_MAP.getOrDefault(dto.getSessionId(),false))
                .map(x->ChatEventVO.builder().eventData(x).eventType(ChatEventTypeEnum.DATA.getValue()).build())//定义返回的内容
                .concatWith(Flux.just(ChatEventVO.builder().eventType(ChatEventTypeEnum.STOP.getValue()).build()));//定义结束标志
    }
    /**
     * 停止会话
     * @param sessionId
     */
    @Override
    public void stop(String sessionId) {
        SESSION_MAP.put(sessionId,false);
    }
}
